# -*- coding: utf-8 -*-
"""
Created on Tue Jun 16 06:29:09 2020

@author: HP
"""

from os.path import isfile
import numpy as np
import pandas as pd
from itertools import permutations
from numpy.random import default_rng
from increasing_closures import lottery_equivalence, \
    impatience, correlation_neutrality, correlation_aversion, \
    correlation_seeking, stochastic_impatience, stochastic_patience
from garp_NOQ import garp_NOQ, garp_NOQ_find_e

#=============================================================================
#SIMULATE FILE
#=============================================================================
sim_file = 'simulate_imple_um.xlsx'
simN = 10000

if isfile(sim_file):
    sim_df = pd.read_excel(sim_file, sheet_name = 0, index_col = 0)
else:
    cols = ['participantid', 'x', 'y', 'z', 'w', 'px', 'py', 'pz', 'pw']
    sim_df = pd.DataFrame(data = None, columns = cols, index = range(simN * 41))
    sim_df['participantid'] = np.ceil( (sim_df.index + 1.) / 41. ).astype(int)
    sim_df.to_excel(sim_file)
    
#==========================================================================
#MAKE PRICE LIST
#==========================================================================

#Make the 96 price vectors
price_list = [0.5,0.8,1.25,2.]
prices_df = pd.DataFrame( np.nan, columns = ['px','py', 'pz','pw'], index = range(96) )
rowi = 0 #keep track of the row I'm working on
for ones_place in range(4):
    #ones_place marks where the 1 will appear.
    #price_list_active holds different 3-permutations of the price_list.
    for price_list_active in permutations( price_list, 3 ):
        prices_df.iloc[ rowi, ones_place ] = 1.
        prices_df.iloc[ rowi, np.array([0,1,2,3]) != ones_place ] = np.array( price_list_active )
        rowi = rowi + 1

#==========================================================================
#MAKE THE CHOICE DATA
#==========================================================================

e = 0.95
Ilist = [ lottery_equivalence, impatience ]

for pid in range(1,simN+1):
    #See if I've already processed pid
    if pd.isnull( sim_df.loc[41*(pid-1),:] ).any():
        print('working on pid: {0}'.format(pid))
        C = np.zeros( ( 41, 4 ) )
        P = np.concatenate( ( np.array( [[1.,1.,1.,1.]] ), prices_df.sample(n=40, replace=False).values ))        
        rng = default_rng()
        for n in range(41):
            flag = False #Turn to true once choices satisfy imple
            while not flag:
                b_vec = rng.dirichlet([1.,1.,1.,1.]) #Draw shares from uniform distribution
                C[n] = (100. * b_vec) / P[n]
                if garp_NOQ( C[0:n+1,:], P[0:n+1,:], Ilist, e ):
                    flag = True
        #The pid data is created
        #Write to the sim_df
        sim_df.loc[ sim_df['participantid'] == pid, ['x','y','z','w'] ] = C
        sim_df.loc[ sim_df['participantid'] == pid, ['px','py','pz','pw'] ] = P
        sim_df.to_excel( sim_file )

#==========================================================================
#RESULTS FILE
#==========================================================================

imple_list = [impatience, lottery_equivalence]
tests_dict = {'def': [], \
              'cn': [correlation_neutrality], \
              'ca': [correlation_aversion], \
              'cs': [correlation_seeking], \
              'si': [stochastic_impatience], \
              'sp': [stochastic_patience]}

columns = list(tests_dict.keys())
pids = sim_df['participantid'].unique()

results_file = 'rp_tests_by_participant_sim_um.xlsx'

#-------------------
#manage results file
#-------------------
if isfile( results_file ):
    results_df = pd.read_excel( results_file, sheet_name = 0, index_col = 0 )
else:
    #Make a new results df
    results_df = pd.DataFrame( np.nan, index = pids, columns = columns )
    results_df.index.name = 'participantid'

#==========================================================================
#RUN TESTS
#==========================================================================
    
counter = 0
for pid in pids:
    if pd.isnull( results_df.loc[pid,:] ).any():
        print('current id: {0}'.format(pid))
        counter = counter + 1
        cur_df = sim_df[ sim_df['participantid'] == pid ]
        C = cur_df[['x','y','z','w']].to_numpy()
        P = cur_df[['px','py','pz','pw']].to_numpy()
        for test_name in tests_dict:
            #Run the test if we don't already have the results.
            e = garp_NOQ_find_e(C,P,Ilist= tests_dict[test_name])
            results_df.loc[pid,test_name] = e
        if counter == 100:
            print('saving results')
            results_df.to_excel(results_file)
            counter = 0
if counter:
    results_df.to_excel(results_file)
    counter = 0

#==========================================================================
#SUMMARIZE RESULTS
#==========================================================================

summary_file = 'simulate_imple_um_summary.xlsx'

#columns and row names
escores = [100,99,95,90]

var_names = list(tests_dict.keys())

summary_df = pd.DataFrame(data=np.nan, index = escores, columns = var_names)

for escore in escores:

    elevel_df = results_df.copy()
    elevel_df[elevel_df < (escore / 100.)] = 0
    elevel_df[elevel_df >= (escore / 100.)] = 1
    summary_df.loc[ escore, : ] = elevel_df.mean()
    
summary_df.to_excel(summary_file)  






